from typing import Optional
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.models as models
import lpips
from torchmetrics.image import MultiScaleStructuralSimilarityIndexMeasure

import math


class MSSSIM(nn.Module):
    def __init__(self, window_size=11, size_average=True, channel=1, levels=5, K1=0.01, K2=0.03):
        super(MSSSIM, self).__init__()
        self.window_size = window_size
        self.size_average = size_average
        self.channel = channel
        self.levels = levels
        self.K1 = K1
        self.K2 = K2
        
        # Default weights for 5 levels as per Wang et al.
        self.register_buffer('weights', torch.FloatTensor([0.0448, 0.2856, 0.3001, 0.2363, 0.1333]))
        self.register_buffer('window', self._create_window(window_size, channel))

    def _create_window(self, window_size, channel):
        _1D_window = self._gaussian(window_size, 1.5).unsqueeze(1)
        _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)
        window = _2D_window.expand(channel, 1, window_size, window_size).contiguous()
        return window

    def _gaussian(self, window_size, sigma):
        gauss = torch.Tensor([math.exp(-(x - window_size//2)**2/float(2*sigma**2)) for x in range(window_size)])
        return gauss/gauss.sum()

    def _ssim(self, img1, img2, window, size_average=True):
        mu1 = F.conv2d(img1, window, padding=self.window_size//2, groups=self.channel)
        mu2 = F.conv2d(img2, window, padding=self.window_size//2, groups=self.channel)

        mu1_sq = mu1.pow(2)
        mu2_sq = mu2.pow(2)
        mu1_mu2 = mu1 * mu2

        sigma1_sq = F.conv2d(img1 * img1, window, padding=self.window_size//2, groups=self.channel) - mu1_sq
        sigma2_sq = F.conv2d(img2 * img2, window, padding=self.window_size//2, groups=self.channel) - mu2_sq
        sigma12 = F.conv2d(img1 * img2, window, padding=self.window_size//2, groups=self.channel) - mu1_mu2

        C1 = (self.K1) ** 2
        C2 = (self.K2) ** 2

        ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2))
        cs_map = (2 * sigma12 + C2) / (sigma1_sq + sigma2_sq + C2)

        if size_average:
            return ssim_map.mean(), cs_map.mean()
        else:
            return ssim_map.mean(1).mean(1).mean(1), cs_map.mean(1).mean(1).mean(1)

    def forward(self, img1, img2):
        # Ensure inputs are valid
        if img1.size() != img2.size():
            raise ValueError(f"Input sizes don't match: {img1.size()} vs {img2.size()}")
        
        # Check minimum image size for multi-scale
        min_size = min(img1.size(-2), img1.size(-1))
        effective_levels = min(self.levels, int(math.log2(min_size // 4)) + 1)
        
        if effective_levels < 1:
            raise ValueError(f"Image too small for MS-SSIM: {img1.size()}")
        
        weights = self.weights[:effective_levels].to(img1.device)
        
        mssim = []
        mcs = []
        
        # Progressive downsampling - start with original images
        current_img1 = img1
        current_img2 = img2

        for i in range(effective_levels):
            # Compute SSIM and CS at current scale
            ssim_val, cs_val = self._ssim(current_img1, current_img2, self.window, self.size_average)
            
            # Ensure positive values for numerical stability
            ssim_val = torch.clamp(ssim_val, min=1e-8, max=1.0)
            cs_val = torch.clamp(cs_val, min=1e-8, max=1.0)

            mssim.append(ssim_val)
            mcs.append(cs_val)
            
            # Downsample for next iteration (except for last level)
            if i < effective_levels - 1:
                current_img1 = F.avg_pool2d(current_img1, kernel_size=2, stride=2)
                current_img2 = F.avg_pool2d(current_img2, kernel_size=2, stride=2)

        # Convert to tensors
        mssim_tensor = torch.stack(mssim)
        mcs_tensor = torch.stack(mcs)
        
        # Compute MS-SSIM using log-space to avoid numerical issues with torch.prod
        if effective_levels > 1:
            # Use log-space computation: log(prod(x^w)) = sum(w * log(x))
            log_cs_weighted = torch.sum(weights[:-1] * torch.log(mcs_tensor[:-1]))
            cs_weighted = torch.exp(log_cs_weighted)
        else:
            cs_weighted = torch.tensor(1.0, device=img1.device)
        
        # Use full SSIM for the last (coarsest) level
        ssim_weighted = mssim_tensor[-1] ** weights[-1]
        
        ms_ssim = cs_weighted * ssim_weighted
        
        # Final safety check and clamp to valid range
        ms_ssim = torch.clamp(ms_ssim, min=1e-8, max=1.0)
        
        # Additional NaN/Inf check (should be rare now)
        if torch.isnan(ms_ssim).any() or torch.isinf(ms_ssim).any():
            print(f"WARNING: Invalid MS-SSIM value detected. Falling back to SSIM.")
            return torch.clamp(mssim_tensor[-1], min=1e-8, max=1.0)
        
        return ms_ssim


class SimpleSSIM(nn.Module):
    def __init__(self, window_size=7, size_average=True, channel=1, K1=0.01, K2=0.03):
        super(SimpleSSIM, self).__init__()
        self.window_size = window_size
        self.size_average = size_average
        self.channel = channel
        self.K1 = K1
        self.K2 = K2

        self.register_buffer('window', self._create_window(window_size, channel))

    def _create_window(self, window_size, channel):
        _1D_window = self._gaussian(window_size, 1.5).unsqueeze(1)
        _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)
        window = _2D_window.expand(channel, 1, window_size, window_size).contiguous()
        return window

    def _gaussian(self, window_size, sigma):
        gauss = torch.Tensor([math.exp(-(x - window_size//2)**2/float(2*sigma**2)) for x in range(window_size)])
        return gauss/gauss.sum()

    def forward(self, img1, img2):
        # Ensure inputs are valid
        if img1.size() != img2.size():
            raise ValueError(f"Input sizes don't match: {img1.size()} vs {img2.size()}")
        
        # Check for NaN/Inf inputs
        if torch.isnan(img1).any() or torch.isnan(img2).any():
            raise ValueError("NaN detected in SSIM inputs")
        if torch.isinf(img1).any() or torch.isinf(img2).any():
            raise ValueError("Inf detected in SSIM inputs")
        
        # Ensure window is on the correct device
        window = self.window.to(img1.device)
        
        # Compute means
        mu1 = F.conv2d(img1, window, padding=self.window_size//2, groups=self.channel)
        mu2 = F.conv2d(img2, window, padding=self.window_size//2, groups=self.channel)

        mu1_sq = mu1.pow(2)
        mu2_sq = mu2.pow(2)
        mu1_mu2 = mu1 * mu2

        # Compute variances and covariance
        sigma1_sq = F.conv2d(img1 * img1, window, padding=self.window_size//2, groups=self.channel) - mu1_sq
        sigma2_sq = F.conv2d(img2 * img2, window, padding=self.window_size//2, groups=self.channel) - mu2_sq
        sigma12 = F.conv2d(img1 * img2, window, padding=self.window_size//2, groups=self.channel) - mu1_mu2

        # SSIM constants
        C1 = (self.K1) ** 2
        C2 = (self.K2) ** 2

        # Compute SSIM map
        numerator = (2 * mu1_mu2 + C1) * (2 * sigma12 + C2)
        denominator = (mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)
        
        # Add small epsilon to avoid division by zero
        denominator = torch.clamp(denominator, min=1e-12)
        ssim_map = numerator / denominator
        
        # Clamp SSIM values to valid range
        ssim_map = torch.clamp(ssim_map, min=-1.0, max=1.0)

        if self.size_average:
            ssim_value = ssim_map.mean()
        else:
            ssim_value = ssim_map.mean(1).mean(1).mean(1)
        
        # Final validity check
        if torch.isnan(ssim_value).any() or torch.isinf(ssim_value).any():
            print(f"WARNING: Invalid SSIM value detected: {ssim_value}")
            return torch.tensor(1.0, device=img1.device)
        
        return ssim_value


class SimplerPerceptualLoss(nn.Module):
    def __init__(self, net='vgg'):
        """
        使用LPIPS的感知损失，适用于单通道图像
        
        参数:
            net: 使用的骨干网络，可选'alex', 'vgg', 'squeeze'
        """
        super(SimplerPerceptualLoss, self).__init__()
        self.loss_fn = lpips.LPIPS(net=net)
        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        if self.device == 'cuda':  # 如果可用GPU，则将模型移动到GPU
            self.loss_fn.cuda()
        for param in self.loss_fn.parameters():
            param.requires_grad = False

    def forward(self, x, y):
        """
        计算LPIPS感知损失
        
        参数:
            x: 预测图像 (B, 1, H, W)
            y: 目标图像 (B, 1, H, W)
        """
        # 检查输入有效性
        if torch.isnan(x).any():
            raise ValueError(f"NaN detected in perceptual loss input x: {x}")
        if torch.isnan(y).any():
            raise ValueError(f"NaN detected in perceptual loss input y: {y}")
        if torch.isinf(x).any():
            raise ValueError(f"Inf detected in perceptual loss input x: {x}")
        if torch.isinf(y).any():
            raise ValueError(f"Inf detected in perceptual loss input y: {y}")
        
        if x.device.type != self.device:
            x = x.to(self.device)
            y = y.to(self.device)
        
        # 归一化到[-1,1]范围，添加小的epsilon避免边界值导致的数值问题
        eps = 1e-8
        x = torch.clamp(x, -1.0 + eps, 1.0 - eps)
        y = torch.clamp(y, -1.0 + eps, 1.0 - eps)
        
        # 单通道转三通道
        if x.size(1) == 1:
            x = x.repeat(1, 3, 1, 1)
            y = y.repeat(1, 3, 1, 1)
        
        loss = self.loss_fn(x, y)
        
        # 检查输出有效性
        if torch.isnan(loss).any():
            raise ValueError(f"NaN detected in LPIPS output: {loss}")
        if torch.isinf(loss).any():
            raise ValueError(f"Inf detected in LPIPS output: {loss}")
        
        return loss.mean()


class FocalLoss(nn.Module):
    """专门处理小目标的Focal Loss"""
    def __init__(self, alpha=1, gamma=2, reduce=True):
        super(FocalLoss, self).__init__()
        self.alpha = alpha
        self.gamma = gamma
        self.reduce = reduce

    def forward(self, inputs, targets):
        ce_loss = F.mse_loss(inputs, targets, reduction='none')
        pt = torch.exp(-ce_loss)
        focal_loss = self.alpha * (1-pt)**self.gamma * ce_loss
        
        if self.reduce:
            return torch.mean(focal_loss)
        else:
            return focal_loss


# 综合损失
class ComprehensiveLoss(nn.Module):
    def __init__(self, data_range=1.0, lambda_perc=0.0, lambda_msssim=0., lambda_l2=1.0, lambda_l1=0.0, lambda_focal=0.0):
        super(ComprehensiveLoss, self).__init__()
        self.data_range = data_range
        self.lambda_perc = lambda_perc
        self.lambda_msssim = lambda_msssim
        self.lambda_l2 = lambda_l2
        self.lambda_l1 = lambda_l1
        self.lambda_focal = lambda_focal

        self.perceptual = SimplerPerceptualLoss(net='vgg')  # 感知损失
        # self.msssim = SimpleSSIM(window_size=11, size_average=True, channel=1)  # MS-SSIM损失
        self.msssim = MultiScaleStructuralSimilarityIndexMeasure(gaussian_kernel=False, data_range=data_range, kernel_size=7)
        self.focal_loss = FocalLoss(alpha=1, gamma=2, reduce=True)  # Focal Loss用于小目标检测

    def forward(self, x, y, weight_map: Optional[torch.Tensor] = None):
        """
        计算感知、MS-SSIM、L2、L1的组合损失
        输入:
            x, y: 输入图像, shape为(batch_size, 1, 256, 128)，值范围为[-1,1]
            weight_map: 可选的权重图, shape (batch_size, 1, H, W) or (H, W)
                        用于对L1和L2损失进行空间加权。
        """
        # 检查输入有效性
        if torch.isnan(x).any():
            raise ValueError(f"NaN detected in loss input x: {x}")
        if torch.isnan(y).any():
            raise ValueError(f"NaN detected in loss input y: {y}")
        if torch.isinf(x).any():
            raise ValueError(f"Inf detected in loss input x: {x}")
        if torch.isinf(y).any():
            raise ValueError(f"Inf detected in loss input y: {y}")
        
        # 添加小的epsilon避免边界值
        eps = 1e-7
        x_norm = torch.clamp(x, -1.0 + eps, 1.0 - eps)
        y_norm = torch.clamp(y, -1.0 + eps, 1.0 - eps)
        
        # 计算MS-SSIM
        # MSSSIM期望输入范围在[0,1]，所以做一个变换
        x_msssim = (x_norm + 1) / 2.0
        y_msssim = (y_norm + 1) / 2.0
        x_msssim = torch.clamp(x_msssim, eps, 1.0 - eps)
        y_msssim = torch.clamp(y_msssim, eps, 1.0 - eps)
        
        ms_ssim_loss = 1 - self.msssim(x_msssim, y_msssim)
        
        if torch.isnan(ms_ssim_loss).any():
            raise ValueError(f"NaN detected in MS-SSIM loss: {ms_ssim_loss}")
        
        # 计算感知损失
        if self.lambda_perc > 0:
            perceptual_loss = self.perceptual(x_norm, y_norm)
        else:
            perceptual_loss = torch.tensor(0.0, device=x.device)
        
        # 计算L2损失
        if weight_map is not None:
            pixel_wise_l2 = F.mse_loss(x, y, reduction='none')
            l2_loss = (pixel_wise_l2 * weight_map).mean()
        else:
            l2_loss = F.mse_loss(x, y)

        # 计算L1损失
        if weight_map is not None:
            pixel_wise_l1 = F.l1_loss(x, y, reduction='none')
            l1_loss = (pixel_wise_l1 * weight_map).mean()
        else:
            l1_loss = F.l1_loss(x, y)
        
        # 检查各项损失的有效性
        if torch.isnan(l1_loss).any():
            raise ValueError(f"NaN detected in L1 loss: {l1_loss}")
        if torch.isnan(l2_loss).any():
            raise ValueError(f"NaN detected in L2 loss: {l2_loss}")
        
        focal_loss = self.focal_loss(x, y)

         # 组合损失
        total_loss = (
            (self.lambda_perc * perceptual_loss if self.lambda_perc > 0 else 0) + 
            (self.lambda_msssim * ms_ssim_loss if self.lambda_msssim > 0 else 0) + 
            (self.lambda_l1 * l1_loss if self.lambda_l1 > 0 else 0) +
            (self.lambda_l2 * l2_loss if self.lambda_l2 > 0 else 0) + 
            (self.lambda_focal * focal_loss if self.lambda_focal > 0 else 0)
        )
        
        # 最终检查
        if torch.isnan(total_loss).any():
            raise ValueError(f"NaN detected in total loss: {total_loss}")
        if torch.isinf(total_loss).any():
            raise ValueError(f"Inf detected in total loss: {total_loss}")
        
        return {"loss": total_loss, "perceptual_loss": perceptual_loss, "ms_ssim_loss": ms_ssim_loss, "l1_loss": l1_loss, "l2_loss": l2_loss, "focal_loss": focal_loss}
